import networkx as nx
import matplotlib.pyplot as plt
from collections import deque
from pulp import *
from tools import generate_gas_network, visualize
from matplotlib.font_manager import FontProperties
import pulp as pulp
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

def calculate_users(G):
    sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']
    users = 0
    users_list = []
    queue = deque(sources)
    visited = set(sources)
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("user_weight"):
            users += G.nodes[node]["user_weight"]
            users_list.append(node)
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return users, users_list

def calculate_leaks(G):
    sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']
    leaks = 0
    leaks_list = []
    queue = deque(sources)
    visited = set(sources)
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("leak_weight"):
            leaks += G.nodes[node]["leak_weight"]
            leaks_list.append(node)
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return leaks, leaks_list

if __name__ == "__main__":
    show = False
    #生成图
    G, pos = generate_gas_network()

    if show:
        nx.draw_networkx(G, pos)
        plt.show()

    #predecessors
    #successors
    #generate_cost_flow_network(G)

    #节点
    leaks = [n for n, attr in G.nodes(data=True) if attr['type'] == 'leak']
    sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']
    valves_ori = [n for n, attr in G.nodes(data=True) if attr['type'] == 'valve']
    users = [n for n, attr in G.nodes(data=True) if attr['type'] == 'user']

    #找到能到达的阀门
    valves = []
    queue = deque(sources)
    visited = set(sources)
    while queue:
        node = queue.popleft()
        if node in valves_ori:
            valves.append(node)
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    
    #每个泄漏点的上游最近的阀门
    valves_cut = []
    queue = deque(leaks)
    visited = set(leaks)
    while queue:
        node = queue.popleft()
        if node in valves and node not in valves_cut:
            valves_cut.append(node)
        else:
            for neighbor in G.predecessors(node):
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append(neighbor)

    print (valves_cut)
    
    G5 = G.copy()
    for u in G.nodes():
        if u in valves_cut:
            for v in G.successors(u):
                G5.remove_edge(u, v)
            G5.remove_node(u)
    leaks_ori, leaks_ori_list = calculate_leaks(G)
    leaks_cut, leaks_cut_list = calculate_leaks(G5)
    users_ori, users_ori_list = calculate_users(G)
    users_cut, users_cut_list = calculate_users(G5)
    print ("阻止泄漏点：", leaks_ori, "-", leaks_cut, "=", leaks_ori - leaks_cut, list(set(leaks_ori_list) - set(leaks_cut_list)))
    print ("影响用户数：", users_ori, "-", users_cut, "=", users_ori - users_cut, list(set(users_ori_list) - set(users_cut_list)))
    visualize(G, pos, valves_cut, list(set(users_ori_list) - set(users_cut_list)))
    